package com.xiao.xin.core.service.Impl;

import com.alibaba.fastjson2.JSON;
import com.xiao.xin.core.common.Constants;
import com.xiao.xin.core.domain.ComfyuiTask;
import com.xiao.xin.core.domain.MessageBase;
import com.xiao.xin.core.service.*;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.redisson.api.RSemaphore;
import org.redisson.api.RedissonClient;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Service;

import java.time.Duration;
import java.time.temporal.ChronoUnit;
import java.util.HashMap;
import java.util.List;
import java.util.function.Function;

@Slf4j
@Service
@RequiredArgsConstructor
public class ComfyuiMessageServiceImpl implements ComfyuiMessageService {

    private final WsNoticeService wsNoticeService;
    private final RedisService redisService;
    private final RedissonClient redissonClient;
    private final StringRedisTemplate redisTemplate;


    @Override
    public void handleMessage(MessageBase message) {
        if ("progress".equals(message.getType())) {
            handleProgress(message);
        } else if ("executed".equals(message.getType())) {
            handleExecuted(message);
        } else if ("execution_error".equals(message.getType())) {
            handleError(message);
        } else if ("status".equals(message.getType())) {
            handleStatus(message);
        }else if ("start".equals(message.getType())){
            handleStart();
        }
    }

    private void handleStart() {
        redisTemplate.opsForValue().set(Constants.START_TASK,"已开始", Duration.of(10, ChronoUnit.MINUTES));
    }

    private void handleStatus(MessageBase message) {
        HashMap<String, Object> data = message.getData();
        HashMap<String, HashMap> status = (HashMap<String, HashMap>) data.get("status");
        Integer count = Integer.valueOf(status.get("exec_info").get("queue_remaining").toString());
        log.info("ComfyUI队列中的数量:{}", count);
        if (count == 0) {
            RSemaphore semaphore = redissonClient.getSemaphore(Constants.TASK_RUN_SEMAPHORE);
            semaphore.release();
            log.info("归还信号量");
        }
    }

    private final UserFundRecordService userFundRecordService;

    private void handleError(MessageBase message) {
        HashMap<String, Object> data = message.getData();
        HashMap<String, Object> hashMap = new HashMap<>();
        hashMap.put("type", message.getType());
        hashMap.put("exception_message", data.get("exception_message"));

        ComfyuiTask comfyuiTask = redisService.getStartedTask(data.get("prompt_id").toString());

        userFundRecordService.freezeReturn(comfyuiTask.getUserId(), comfyuiTask.getSize());

        wsNoticeService.sendToUser(comfyuiTask.getWsClientId(), JSON.toJSONString(hashMap));
    }

    private final UserResultService userResultService;

    private void handleExecuted(MessageBase message) {
        HashMap<String, Object> data = message.getData();
        HashMap<String, Object> output = (HashMap<String, Object>) data.get("output");
        List<HashMap<String, Object>> images = (List<HashMap<String, Object>>) output.get("images");
        List<String> urls = images.stream().map(new Function<HashMap<String, Object>, String>() {
            @Override
            public String apply(HashMap<String, Object> hashMap) {
                return String.format("http://192.168.100.129:8188/view?filename=%s&type=%s&subfolder=%s", hashMap.get("filename"), hashMap.get("type"), hashMap.get("subfolder"));
            }
        }).toList();
        log.info("处理完工的Urls:{}", urls);


        HashMap<String, Object> hashMap = new HashMap<>();
        hashMap.put("type", "imageResult");
        hashMap.put("urls", urls);


        ComfyuiTask comfyuiTask = redisService.getStartedTask(data.get("prompt_id").toString());
        wsNoticeService.sendToUser(comfyuiTask.getWsClientId(), JSON.toJSONString(hashMap));

        userResultService.insertUserResult(urls, comfyuiTask.getUserId());

    }


    private void handleProgress(MessageBase message) {
        HashMap<String, Object> map = message.getData();
        String promptId = map.get("prompt_id").toString();
        ComfyuiTask clientId = redisService.getStartedTask(promptId);
        map.put("type", "progress");
        wsNoticeService.sendToUser(clientId.getWsClientId(), JSON.toJSONString(map));
    }
}
